﻿using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Linq;
using System.Text;

namespace MachinaAurum.DempsterShafer
{
    public class DempsterShaferEngine
    {
        public IDictionary<int, string> UniversalSet { get; private set; }

        public IEnumerable<ISet<int>> PowerSetValues { get; private set; }
        public IEnumerable<ISet<string>> PowerSetLabels { get; private set; }

        List<Evidence> Evidences = new List<Evidence>();

        public IDictionary<ISet<int>, double> Current { get; private set; }

        public DempsterShaferEngine()
        {
            UniversalSet = new Dictionary<int, string>();
        }

        public void Start()
        {
            UniversalSet = new ReadOnlyDictionary<int, string>(UniversalSet);

            var values = UniversalSet.Select(x => x.Key);
            var labels = UniversalSet.Select(x => x.Value);

            PowerSetValues = GetPowerSet(values).ToArray();
            PowerSetLabels = GetPowerSet(labels).ToArray();
        }

        public void AddState(int value, string label)
        {
            UniversalSet.Add(value, label);
        }

        IEnumerable<ISet<T>> GetPowerSet<T>(IEnumerable<T> input)
        {
            var seed = new List<ISet<T>>() { new SortedSet<T>() } as IEnumerable<ISet<T>>;

            return input.Aggregate(seed, (a, b) => a.Concat(a.Select(x =>
            {
                var newSet = new SortedSet<T>(x);
                newSet.UnionWith(new[] { b });
                return newSet;
            }))).ToArray();
        }

        public Evidence AddEvidences()
        {
            return new Evidence(PowerSetValues);
        }

        public void Accept(Evidence evidence)
        {
            Evidences.Add(evidence);

            if (Evidences.Count > 1)
            {
                var normalizedEvidences = CombineAndNormalize();

                Current = new ReadOnlyDictionary<ISet<int>, double>(normalizedEvidences);
            }
            else
            {
                var currentMasses = PowerSetValues.ToDictionary(x => x, x => evidence.GetMass(x));
                Current = new ReadOnlyDictionary<ISet<int>, double>(currentMasses);
            }
        }

        public double GetCurrentBeliefIn(params int[] state)
        {
            if (state.Length > 0)
            {
                var powerSet = GetPowerSet(state).Where(x => x.Any());
                return powerSet.Sum(x => Current[x]);
            }
            else
            {
                return Current[new SortedSet<int>()];
            }
        }

        private IDictionary<ISet<int>, double> CombineAndNormalize()
        {
            var m1 = Current;
            var m2 = Evidences.Last();

            var d = PowerSetValues.ToDictionary(x => x, x => 0.0);
            var states = new SetDictionary<int, double>(d);

            var comparer = new SetComparer<int>();
            var combinations = from theta1 in m1.Select(x => x.Key).ToArray()
                               from theta2 in m2.States
                               let theta3 = GetIntersection(theta1, theta2)
                               let mass1 = m1[theta1]
                               let mass2 = m2.GetMass(theta2)
                               let comb = new { States = theta3, Mass = mass1 * mass2 }
                               group comb by comparer.GetHashCode(theta3) into grouped
                               select grouped;

            double k = 0.0;

            foreach (var combination in combinations)
            {
                var theta3 = combination.First().States;

                if (theta3 == null)
                {
                    k += combination.Sum(x => x.Mass);
                }
                else
                {
                    states[theta3] = combination.Sum(x => x.Mass);
                }
            }

            if (k > 0)
            {
                foreach (var state in states.ToArray())
                {
                    states[state.Key] = state.Value / (1.0 - k);
                }
            }

            return states;
        }

        private static ISet<int> GetIntersection(ISet<int> theta1, ISet<int> theta2)
        {
            var isTheta1Empty = theta1.Any() == false;
            var isTheta2Empty = theta2.Any() == false;

            if (isTheta1Empty && isTheta2Empty)
            {
                return theta1;
            }
            else if (isTheta1Empty)
            {
                return theta2;
            }
            else if (isTheta2Empty)
            {
                return theta1;
            }

            var intersection = new SortedSet<int>(theta1);
            intersection.IntersectWith(theta2);

            if (intersection.Any() == false)
            {
                return null;
            }

            return intersection;
        }

        public double GetCurrentPlausabilityIn(params int[] states)
        {
            var notStates = UniversalSet.Select(x => x.Key).Except(states).ToArray();
            return 1.0 - GetCurrentBeliefIn(notStates);
        }

        public double GetCurrentDoubtIn(params int[] states)
        {
            var notStates = UniversalSet.Select(x => x.Key).Except(states).ToArray();
            return GetCurrentBeliefIn(notStates);
        }

        public double GetCurrentIgnoranceIn(params int[] states)
        {
            return GetCurrentPlausabilityIn(states) - GetCurrentBeliefIn(states);
        }
    }

    public class SetDictionary<TKey, TValue> : Dictionary<ISet<TKey>, TValue>
    {
        public SetDictionary()
            : base(new SetComparer<TKey>())
        {
        }

        public SetDictionary(IDictionary<ISet<TKey>, TValue> dictionary)
            : base(dictionary, new SetComparer<TKey>())
        {
        }
    }

    public class SetComparer<T> : IEqualityComparer<ISet<T>>
    {
        public bool Equals(ISet<T> x, ISet<T> y)
        {
            return x.SetEquals(y);
        }

        public int GetHashCode(ISet<T> obj)
        {
            if (obj == null)
            {
                return 0;
            }

            var builder = new StringBuilder();

            foreach (var item in obj.OrderBy(x => x))
            {
                builder.Append(item);
            }

            return builder.ToString().GetHashCode();
        }
    }
}
